open Ast
open Bytecode


module StringMap = Map.Make(String)

(* Symbol table: Information about all the names in scope *)
type env = {
    function_index : int StringMap.t; (* Index for each function *)
    global_index   : int StringMap.t; (* "Address" for global variables *)
    local_index    : int StringMap.t; (* FP offset for args, locals *)
  }


(*this is the function to calculate the list inside*)

let rec inenum n x =function
	                  []->[] 
	                    |hd::tl->(n, hd)::(inenum (n+x) x tl)  


(*[(Question,[Q1;Q2]);(QuestionSet,[Qs1;Qs2])]*)

let rec enum n= function
  		[]->[]
		|hd::tl-> let le= (match fst hd with
			  Question->13
			 |QuestionSet->4
			 |Image->2
             |Jump->3) in 
                                  let next=(List.length (snd hd))*le+n 
                                    in (inenum n le (snd (hd)))@(enum next tl)


let rec enumf n= function
  		[]->[]
		|hd::tl-> let le= (match fst hd with
			  Question->(-13)
			 |QuestionSet->(-4)
			 |Image->(-2)
             |Jump ->(-2) ) in 
                                  let next=(List.length (snd hd))*le+n 
                                    in (inenum n le (snd (hd)))@(enum next tl)
                                    

let rec formenum n= function
  		[]->[]
		|hd::tl-> let le= (match fst hd with
			  Question->(-13)
			 |QuestionSet->(-4)
			 |Image->(-2)
             |Jump ->(-2)) in  (n, snd(hd))::(formenum (n+le) tl)                                  
                                    

                                    
                                    
                                    

let arglenfun s e=
                    let le=(match fst e with
				  Question->13
				  |QuestionSet->4
				  |Image->2
                  |Jump ->2) in 
					let arglen=((List.length (snd e))*le) in
					    arglen+s 

let formlenfun s e=
                    let le=(match fst e with
				  Question->13
				  |QuestionSet->4
				  |Image->2
                  |Jump ->2)
					 in  le+s   (*need to check*)


               

let string_map_pairs map pairs=
		List.fold_left ( fun m (i, n)-> StringMap.add n i m) map pairs


let translate (globals, functions)=
   
     let global_indexes =
		string_map_pairs StringMap.empty (enum 0 globals) in

     let function_indexes=string_map_pairs StringMap.empty
				(inenum 1 1 (List.map (function f->f.fname) functions)) in
    
     let translate env fdecl=
      let formals_length= List.fold_left formlenfun 0 fdecl.formals
        and locals_length= List.fold_left arglenfun 0 fdecl.locals  
        and locals_offsets= enum 1 fdecl.locals
        and formals_offsets= formenum (-1) fdecl.formals in

        let env={ env with local_index= string_map_pairs 
			StringMap.empty (locals_offsets@formals_offsets) 
		}  in
      
       let rec expr= function
        | Link (e1, e2) -> (try [LoLkwb ((StringMap.find e1 env.local_index), e2) ]
                                  with Not_found-> try
                                    [GlLkwb((StringMap.find e1 env.local_index), e2)] 
                                    with Not_found->
                                    raise (Failure ("undeclared variable" ^ e1)))
                                    
        |Jup (e1, e2)-> (try [LoJpwb ((StringMap.find e1 env.local_index), e2)]
                                    with Not_found-> try
                                        [GlJpwb ((StringMap.find e1 env.local_index), e2)]
                                    with Not_found-> 
                                    raise (Failure ("undeclared variable" ^ e2)))          
        | Title(e1,e2)->(try [Lowbt ((StringMap.find e1 env.local_index), e2, e1)]
		                          with Not_found-> try
                                        [Glwbt ((StringMap.find e1 env.global_index),e2, e1)]
		                                   with Not_found->
			                                 raise (Failure ("undeclared variable" ^ e1)))

	    |Atype(e1,form,e2)->  let atlist=(List.concat (List.map (function a->
                                                                 try [Pushal ((StringMap.find e1 env.local_index), a)]
                                                                   with Not_found-> try
                                                                      [Pushag ((StringMap.find e1 env.global_index), a)]
                                                                      with Not_found->
                                                                      raise (Failure ("undeclared variable" ^ e1 ))) e2))
			     and tp=( match form with
					     Single->1
					     |Multiple->2
					     |Drop->3
				         )
				        in( try [Lowba ((StringMap.find e1 env.local_index),tp)]@atlist
					with Not_found-> try
					         [Glwba ((StringMap.find e1 env.global_index),tp)]@atlist
				        with Not_found->
					  raise (Failure ("undeclared variable" ^ e1)))
        |Text(e1,e2)->( try [Lotwbt ((StringMap.find e1 env.local_index), int_of_string e2)]
				    with Not_found-> try
					[Gltwbt ((StringMap.find e1 env.global_index), int_of_string e2)]
				    with Not_found->
					raise (Failure ("undeclared variable"^ e1))					
				)

	   |Addquestion(e1,e2)->let len=(List.length e2)
			     and qlist=(List.map (function a-> try (Loquetrace (StringMap.find a env.local_index))
							       with Not_found -> 
								   raise (Failure ("undecluared variable"^ a))
                                                  ) e2)
					in (try ([Loqswb ((StringMap.find e1 env.local_index),len)])@qlist
				            with Not_found-> 
						raise (Failure ("undeclared variable"^ e1))
						)
    | Call (fname, actuals) -> (try
	  (List.concat (List.map (function a -> try [Lfp (StringMap.find a env.local_index)]
                                                with Not_found-> try
                                                [Lod (StringMap.find a env.global_index)]
                                                with Not_found-> 
                                                raise (Failure ("undeclared variable" ^ a))
                                            ) (List.rev actuals))) @
	  [Jsr (StringMap.find fname env.function_index) ]   
        with Not_found -> raise (Failure ("undefined function " ^ fname)))
        
    
    |Display(e1)-> (List.concat (List.map ( function a-> try [Loshow (StringMap.find a env.local_index)]
					                                    with Not_found-> 
					                                       raise (Failure ("undeclared variable"^ a))) e1))

   in let rec stmt= function
        Expr e->expr e
	|Block s1-> List.concat (List.map stmt s1)

   in [Ent locals_length]@ 
     stmt (Block fdecl.body)@
	[Rts formals_length] 
 
   in let env = { function_index = function_indexes;
		 global_index = global_indexes;
		 local_index = StringMap.empty } in

  (* Code executed to start the program: Jsr main; halt *)
  let entry_function = try
    [Jsr (StringMap.find "main" function_indexes); Hlt]
  with Not_found -> raise (Failure ("no \"main\" function"))
  in
    
  (* Compile the functions *)
  let func_bodies = entry_function :: List.map (translate env) functions in

  (* Calculate function entry points by adding their lengths *)
  let (fun_offset_list, _) = List.fold_left
      (fun (l,i) f -> (i :: l, (i + List.length f))) ([],0) func_bodies in
  let func_offset = Array.of_list (List.rev fun_offset_list) in

  { num_globals =List.fold_left arglenfun 0 globals;
    (* Concatenate the compiled functions and replace the function
       indexes in Jsr statements with PC values *)
    text = Array.of_list (List.map (function
	Jsr i when i > 0 -> Jsr func_offset.(i)
      | _ as s -> s) (List.concat func_bodies))
  }
            







